import os
import pandas as pd
from pathlib import Path
import dolphindb as ddb
import traceback

class ZCEImporter:
    """郑州商品交易所数据导入器"""
    
    def __init__(self):
        self.source_path = Path("F:/资料/历史数据/郑商所/data")
        self.session = None
        self.upserter = None
        
    def connect_db(self):
        """连接DolphinDB数据库"""
        self.session = ddb.session()
        self.session.connect("localhost", 8848, "admin", "123456")
        self.upserter = ddb.TableUpserter(
            dbPath="dfs://futures", 
            tableName="daily_all", 
            ddbSession=self.session,
            keyColNames=["date", "code", "exchange_code"]
        )

    def clean_data(self, df):
        """清理数据中的制表符和空格"""
        for col in df.columns:
            if df[col].dtype == 'object':
                df[col] = df[col].astype(str).apply(lambda x: x.strip().strip('\t'))
        return df

    def clean_number(self, x):
        """清理数值字符串，移除千位分隔符并保留两位小数"""
        try:
            if isinstance(x, str):
                return round(float(x.replace(',', '')), 2)
            return round(float(x), 2)
        except:
            return None

    def is_option_code(self, code):
        """判断是否为期权合约"""
        try:
            code = str(code).upper()
            option_prefixes = ['CO', 'PO', 'ZO', 'MO', 'SO', 'RO']
            return any(code.startswith(prefix) for prefix in option_prefixes) or 'O' in code
        except:
            return False

    def _get_product_name(self, code):
        """从合约代码提取品种名称"""
        try:
            code = str(code)
            product_names = {
                'CF': '棉花',
                'SR': '白糖',
                'TA': 'PTA',
                'OI': '菜油',
                'RM': '菜粕',
                'MA': '甲醇',
                'FG': '玻璃',
                'RS': '菜籽',
                'RI': '早籼稻',
                'WH': '强麦',
                'JR': '粳稻',
                'LR': '晚籼稻',
                'PM': '普麦',
                'SF': '硅铁',
                'SM': '锰硅',
                'ZC': '动力煤',
                'AP': '苹果',
                'CJ': '红枣'
            }
            product_code = ''.join(filter(str.isalpha, code[:2].upper()))
            return product_names.get(product_code, product_code)
        except:
            return code

    def process_file(self, file_path):
        """处理单个文件"""
        try:
            # 尝试不同的编码读取文件
            for encoding in ['gbk', 'gb2312', 'gb18030', 'utf-8']:
                try:
                    # 读取文件头
                    with open(file_path, 'r', encoding=encoding) as f:
                        _ = f.readline()  # 跳过标题行
                        header = f.readline().strip()
                    
                    # 处理列名
                    columns = [col.strip().strip('\t') for col in header.split('|')]
                    
                    # 读取数据
                    df = pd.read_csv(file_path, 
                                   encoding=encoding,
                                   sep='|',
                                   skiprows=2,
                                   names=columns,
                                   skipinitialspace=True,
                                   na_values=[''])
                    break
                except:
                    continue
            else:
                print(f"无法读取文件 {file_path}")
                return None

            # 基础数据清理
            df = df.dropna(how='all')
            if df.columns[-1].strip() == '':
                df = df.iloc[:, :-1]
            df = self.clean_data(df)

            # 标准化列名
            rename_dict = {
                '交易日期': 'date',
                '品种月份': 'code',
                '今收盘': 'close',
                '今开盘': 'open',
                '最高价': 'high',
                '最低价': 'low',
                '昨结算': 'pre_close',
                '成交量(手)': 'volume',
                '成交额(万元)': 'amount'
            }

            # 重命名列
            df = df.rename(columns={k: v for k, v in rename_dict.items() if k in df.columns})

            # 数值转换
            numeric_cols = ['open', 'high', 'low', 'close', 'pre_close']
            for col in numeric_cols:
                if col in df.columns:
                    df[col] = df[col].apply(self.clean_number)

            # 处理成交量和成交额
            if 'volume' in df.columns:
                df['volume'] = df['volume'].apply(self.clean_number)
            
            if 'amount' in df.columns:
                df['amount'] = df['amount'].apply(lambda x: 
                    round(self.clean_number(x) * 10000, 2) if x is not None else None)

            # 处理日期格式
            df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')

            # 添加交易所信息
            df['exchange_code'] = 'ZCE'
            df['exchange_name'] = '郑商所'

            # 提取品种名称并过滤期权
            df['name'] = df['code'].apply(self._get_product_name)
            df = df[~df['code'].apply(self.is_option_code)]

            # 选择和排序最终列
            final_columns = ['date', 'code', 'name', 'exchange_code', 'exchange_name',
                           'high', 'low', 'open', 'close', 'pre_close', 'volume', 'amount']

            return df[final_columns]

        except Exception as e:
            print(f"处理文件失败 {file_path}: {str(e)}")
            print(f"错误类型: {type(e).__name__}")
            print(f"详细错误信息: {traceback.format_exc()}")
            return None

    def import_to_db(self, df):
        """导入数据到数据库"""
        try:
            self.upserter.upsert(df)
            print(f"成功处理 {len(df)} 条数据")
        except Exception as e:
            print(f"导入数据失败: {str(e)}")
            try:
                print("尝试重新连接...")
                self.connect_db()
                self.upserter.upsert(df)
                print(f"重连后成功处理 {len(df)} 条数据")
            except Exception as e2:
                print(f"重试失败: {str(e2)}")

    def process_all_files(self):
        """处理所有文件"""
        try:
            self.connect_db()
            
            for file_path in self.source_path.glob('*.txt'):
                print(f"\n处理文件: {file_path.name}")
                df = self.process_file(file_path)
                
                if df is not None and not df.empty:
                    print(f"处理得到 {len(df)} 行数据")
                    self.import_to_db(df)
                    
        except Exception as e:
            print(f"处理过程出错: {str(e)}")
            raise
        finally:
            if self.session:
                self.session.close()

def main():
    """主函数"""
    importer = ZCEImporter()
    importer.process_all_files()

if __name__ == "__main__":
    main() 